# Copyright (c) 2020 cgsdfc
#
# This software is released under the MIT License.
# https://opensource.org/licenses/MIT
"""
Generate a Python.asdl file from the Python-ast.h header.
"""

from pathlib import Path
from typing import Sequence, Tuple

from pycparser.c_ast import (FileAST, Typedef, Decl, Enum, Struct, TypeDecl,
                             Union, Enumerator, EnumeratorList, Node, PtrDecl,
                             IdentifierType)
from pycparser import parse_file as ParseFile

from astlib.asdl import (Module, Type, Constructor, Field, Sum, Product,
                         builtin_types as ASDL_BUILTIN_TYPES, check as CheckAsdlModule)

from astlib.asdl_unparse import UnparseAsdlModule, ASDL_MODULE_ATTR_COMMENT
from astlib.asdl_index import AsdlIndex
from astlib.license_const import LICENSE_BANG

ASDL_INDEX = AsdlIndex()


class NameUtils:
    @classmethod
    def GetTypeName(cls, type) -> str:
        if isinstance(type, IdentifierType):
            return type.names[0]
        if isinstance(type, PtrDecl):
            assert isinstance(type.type, TypeDecl)
            assert isinstance(type.type.type, IdentifierType)
            return cls.GetTypeName(type.type.type)
        if isinstance(type, TypeDecl):
            assert isinstance(type.type, IdentifierType)
        return cls.GetTypeName(type.type)

    @classmethod
    def Demangle(cls, c_name):
        if c_name in ASDL_BUILTIN_TYPES:
            return c_name
        if c_name.endswith('_ty'):
            return c_name[:-3]
        if c_name.startswith('_'):
            return c_name[1:]
        return c_name


class SumEnum:
    @classmethod
    def IsSumEnum(cls, typedef: Typedef):
        """
        Things like:
        `typedef enum _boolop { And=1, Or=2 } boolop_ty;`
        """
        return (isinstance(typedef, Typedef)
                and isinstance(typedef.type, TypeDecl)
                and isinstance(typedef.type.type, Enum))

    @classmethod
    def GetEnumerators(cls, enum: Enum) -> EnumeratorList:
        return enum.values

    @classmethod
    def CreateConstructor(cls, enumerator: Enumerator):
        return Constructor(name=enumerator.name)

    @classmethod
    def CreateSum(cls, enumerator_list: EnumeratorList):
        return Sum(types=[
            cls.CreateConstructor(e) for e in enumerator_list.enumerators
        ])

    @classmethod
    def CreateType(cls, typedef: Typedef):
        name = NameUtils.Demangle(typedef.name)
        enum = typedef.type.type
        return Type(name=name, value=cls.CreateSum(cls.GetEnumerators(enum)))


class Member:

    KNOWN_ATTRIBUTE_NAMES = frozenset([
        'lineno',
        'col_offset',
        'end_lineno',
        'end_col_offset',
    ])

    SEQUENCE_NAMES = frozenset(['asdl_seq', 'asdl_int_seq'])

    @classmethod
    def IsSequence(cls, name):
        return name in cls.SEQUENCE_NAMES

    @classmethod
    def IsMember(cls, decl: Decl):
        return isinstance(decl, Decl) and decl.name != None

    @classmethod
    def GetMemberName(cls, decl: Decl):
        assert cls.IsMember(decl)
        return decl.name

    @classmethod
    def GetMemberTypeName(cls, decl: Decl):
        # assert cls.IsMember(decl)
        c_name = NameUtils.GetTypeName(decl)
        name = NameUtils.Demangle(c_name)
        return name

    @classmethod
    def IsAttribute(cls, decl: Decl):
        return (cls.IsMember(decl)
                and cls.GetMemberName(decl) in cls.KNOWN_ATTRIBUTE_NAMES)

    @classmethod
    def CreateFieldRaw(cls, decl: Decl) -> Field:
        return Field(type=cls.GetMemberTypeName(decl),
                     name=cls.GetMemberName(decl))

    @classmethod
    def CreateField(cls, name, decl: Decl) -> Tuple[Field, bool]:
        f = cls.CreateFieldRaw(decl)
        # When a member is a sequence type, the C code cannot tell us its element
        # type, so we need to look for the complete field from the asdl_index.
        if cls.IsSequence(f.type):
            f = ASDL_INDEX.FindField(name, f.name)

        return f, cls.IsAttribute(decl)


class StructBase:
    @classmethod
    def IsStruct(cls, decl: Decl):
        return isinstance(decl, Decl) and isinstance(decl.type, Struct)

    @classmethod
    def GetStructName(cls, decl: Decl):
        assert cls.IsStruct(decl)
        return decl.type.name

    @classmethod
    def GetMembers(cls, decl: Decl) -> Sequence[Decl]:
        return decl.type.decls

    @classmethod
    def CreateFieldsNoAttrs(cls, name, members: Sequence[Decl]):
        """This don't separate fields from attributes,
        which is useful when client knows there is no attrs."""
        return list([
            result[0] for result in map(
                lambda decl: Member.CreateField(name, decl), members)
        ])

    @classmethod
    def CreateFields(cls, name, members: Sequence[Decl]):
        fields = []
        attrs = []

        def CreateField(decl):
            return Member.CreateField(name, decl)

        for f, is_attr in map(CreateField, members):
            (attrs if is_attr else fields).append(f)
        return fields, attrs

    @classmethod
    def HasUnionMember(cls, decl: Decl):
        return any(
            isinstance(m.type.type, Union) for m in cls.GetMembers(decl))

    @classmethod
    def HasEnumMember(cls, decl: Decl):
        return any(isinstance(m.type.type, Enum) for m in cls.GetMembers(decl))


class ProductStruct(StructBase):
    """
    Things like:
    ```
    struct _withitem {
        expr_ty context_expr;
        expr_ty optional_vars;
    };
    ```
    """
    @classmethod
    def IsProductStruct(cls, decl: Decl):
        return cls.IsStruct(decl) and not cls.HasUnionMember(decl)

    @classmethod
    def CreateProduct(cls, name, members: Sequence[Decl]):
        fields, attributes = cls.CreateFields(name, members)
        return Product(fields=fields, attributes=attributes)

    @classmethod
    def CreateType(cls, decl: Decl):
        assert cls.IsProductStruct(decl)
        name = cls.GetStructName(decl)
        demangled_name = NameUtils.Demangle(name)
        return Type(name=demangled_name,
                    value=cls.CreateProduct(demangled_name,
                                            cls.GetMembers(decl)))


class ConstructorStruct(StructBase):
    """
    Things like:
    ```
    struct _mod {
        enum _mod_kind kind;
        union {
            struct {
                asdl_seq *body;
            } Module;
            
            struct {
                asdl_seq *body;
            } Interactive;
    ```
    """

    ENUM_MEMBER_INDEX = 0
    UNION_MEMBER_INDEX = 1

    @classmethod
    def GetEnumMember(cls, decl: Decl):
        return decl.type.decls[cls.ENUM_MEMBER_INDEX]

    @classmethod
    def GetUnionMember(cls, decl: Decl):
        return decl.type.decls[cls.UNION_MEMBER_INDEX]

    @classmethod
    def IsConstructorStruct(cls, decl: Decl):
        return (cls.IsStruct(decl)
                and isinstance(cls.GetEnumMember(decl).type.type, Enum)
                and isinstance(cls.GetUnionMember(decl).type.type, Union))

    @classmethod
    def IsUnnamedInnerStruct(cls, decl: Decl):
        return cls.IsStruct(decl) and decl.type.type.name is None

    @classmethod
    def CreateConstructor(cls, decl: Decl):
        name = NameUtils.Demangle(decl.name)
        fields = cls.CreateFieldsNoAttrs(name, cls.GetMembers(decl.type))
        return Constructor(name=name, fields=fields)

    @classmethod
    def CreateType(cls, decl: Decl):
        assert cls.IsConstructorStruct(decl)
        name = cls.GetStructName(decl)
        demanged_name = NameUtils.Demangle(name)

        attributes = []
        for m in cls.GetMembers(decl):
            if Member.IsAttribute(m):
                attributes.append(Member.CreateField(demanged_name, m)[0])

        types = [
            cls.CreateConstructor(m)
            for m in cls.GetMembers(cls.GetUnionMember(decl).type)
        ]

        sum = Sum(types=types, attributes=attributes)
        return Type(name=demanged_name, value=sum)


class ModuleGenerator:

    MODULE_NAME = 'Python'

    @classmethod
    def MaybeCreateType(cls, node: Node):
        if SumEnum.IsSumEnum(node):
            return SumEnum.CreateType(node)
        if ProductStruct.IsProductStruct(node):
            return ProductStruct.CreateType(node)
        if ConstructorStruct.IsConstructorStruct(node):
            return ConstructorStruct.CreateType(node)
        return None

    @classmethod
    def CreateModule(cls, file_ast: FileAST) -> Module:
        dfns = list(filter(None, map(cls.MaybeCreateType, file_ast.ext)))
        # TODO: don't setattr()
        mod = Module(name=cls.MODULE_NAME, dfns=dfns)
        setattr(mod, ASDL_MODULE_ATTR_COMMENT, LICENSE_BANG)
        assert CheckAsdlModule(mod)
        return mod

    @classmethod
    def GenerateAsdlFromHeader(cls, hdr: Path, output: Path):
        file_ast: FileAST = ParseFile(filename=str(hdr))
        mod: Module = cls.CreateModule(file_ast)
        UnparseAsdlModule(mod, output)


if __name__ == "__main__":
    from argparse import ArgumentParser
    p = ArgumentParser()
    p.add_argument('-i', dest='input', help='Input pp_out.h')
    p.add_argument('-o', dest='output', help='Output Python.asdl')
    args = p.parse_args()
    ModuleGenerator.GenerateAsdlFromHeader(hdr=Path(args.input),
                                           output=Path(args.output))
